#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
化合物数据库查询模块
支持从公共化合物资源库（PubChem, ChemSpider等）查询和下载化合物数据
"""

import json
import time
import requests
from typing import Dict, List, Optional, Tuple
from pathlib import Path
from dataclasses import dataclass, asdict
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class CompoundData:
    """化合物数据结构"""
    cid: Optional[int] = None  # PubChem CID
    cas_number: Optional[str] = None
    name: str = ""
    iupac_name: Optional[str] = None
    molecular_formula: str = ""
    molecular_weight: float = 0.0
    smiles: Optional[str] = None
    inchi: Optional[str] = None
    inchi_key: Optional[str] = None
    
    # 物理性质
    melting_point: Optional[float] = None  # °C
    boiling_point: Optional[float] = None  # °C
    density: Optional[float] = None  # g/cm³
    solubility: Optional[str] = None
    
    # 热力学数据
    enthalpy_formation: Optional[float] = None  # kJ/mol
    gibbs_energy: Optional[float] = None  # kJ/mol
    entropy: Optional[float] = None  # J/(mol·K)
    heat_capacity: Optional[float] = None  # J/(mol·K)
    
    # 安全信息
    hazard_class: Optional[str] = None
    flash_point: Optional[float] = None
    
    # 来源
    data_source: str = "PubChem"
    
    def to_dict(self) -> Dict:
        """转换为字典"""
        return asdict(self)


class CompoundDatabase:
    """化合物数据库查询类"""
    
    # PubChem API基础URL
    PUBCHEM_BASE_URL = "https://pubchem.ncbi.nlm.nih.gov/rest/pug"
    
    # 请求延迟（秒）- 遵守API使用规范
    REQUEST_DELAY = 0.2
    
    def __init__(self, cache_dir: str = "data/compounds"):
        """
        初始化化合物数据库
        
        Args:
            cache_dir: 本地缓存目录
        """
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.cache_file = self.cache_dir / "compound_cache.json"
        self.cache = self._load_cache()
        
    def _load_cache(self) -> Dict:
        """加载本地缓存"""
        if self.cache_file.exists():
            with open(self.cache_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        return {}
    
    def _save_cache(self):
        """保存缓存"""
        with open(self.cache_file, 'w', encoding='utf-8') as f:
            json.dump(self.cache, f, ensure_ascii=False, indent=2)
    
    def search_by_name(self, name: str, use_cache: bool = True) -> Optional[CompoundData]:
        """
        通过化合物名称搜索
        
        Args:
            name: 化合物名称
            use_cache: 是否使用缓存
            
        Returns:
            化合物数据
        """
        # 检查缓存
        cache_key = f"name:{name}"
        if use_cache and cache_key in self.cache:
            logger.info(f"从缓存加载: {name}")
            return CompoundData(**self.cache[cache_key])
        
        logger.info(f"从PubChem查询: {name}")
        
        try:
            # 1. 通过名称获取CID
            url = f"{self.PUBCHEM_BASE_URL}/compound/name/{name}/cids/JSON"
            response = requests.get(url, timeout=10)
            time.sleep(self.REQUEST_DELAY)
            
            if response.status_code != 200:
                logger.warning(f"未找到化合物: {name}")
                return None
            
            data = response.json()
            cid = data['IdentifierList']['CID'][0]
            
            # 2. 通过CID获取详细信息
            compound = self.get_by_cid(cid, use_cache=False)
            
            # 保存到缓存
            if compound:
                self.cache[cache_key] = compound.to_dict()
                self._save_cache()
            
            return compound
            
        except Exception as e:
            logger.error(f"查询失败 {name}: {str(e)}")
            return None
    
    def search_by_cas(self, cas_number: str, use_cache: bool = True) -> Optional[CompoundData]:
        """
        通过CAS号搜索
        
        Args:
            cas_number: CAS编号
            use_cache: 是否使用缓存
            
        Returns:
            化合物数据
        """
        cache_key = f"cas:{cas_number}"
        if use_cache and cache_key in self.cache:
            logger.info(f"从缓存加载 CAS: {cas_number}")
            return CompoundData(**self.cache[cache_key])
        
        logger.info(f"从PubChem查询 CAS: {cas_number}")
        
        try:
            # 通过CAS号获取CID
            url = f"{self.PUBCHEM_BASE_URL}/compound/name/{cas_number}/cids/JSON"
            response = requests.get(url, timeout=10)
            time.sleep(self.REQUEST_DELAY)
            
            if response.status_code != 200:
                logger.warning(f"未找到 CAS: {cas_number}")
                return None
            
            data = response.json()
            cid = data['IdentifierList']['CID'][0]
            
            compound = self.get_by_cid(cid, use_cache=False)
            
            if compound:
                compound.cas_number = cas_number
                self.cache[cache_key] = compound.to_dict()
                self._save_cache()
            
            return compound
            
        except Exception as e:
            logger.error(f"查询失败 CAS {cas_number}: {str(e)}")
            return None
    
    def get_by_cid(self, cid: int, use_cache: bool = True) -> Optional[CompoundData]:
        """
        通过PubChem CID获取化合物信息
        
        Args:
            cid: PubChem Compound ID
            use_cache: 是否使用缓存
            
        Returns:
            化合物数据
        """
        cache_key = f"cid:{cid}"
        if use_cache and cache_key in self.cache:
            logger.info(f"从缓存加载 CID: {cid}")
            return CompoundData(**self.cache[cache_key])
        
        logger.info(f"从PubChem查询 CID: {cid}")
        
        try:
            # 获取化合物属性
            url = f"{self.PUBCHEM_BASE_URL}/compound/cid/{cid}/property/MolecularFormula,MolecularWeight,IUPACName,InChI,InChIKey,CanonicalSMILES/JSON"
            response = requests.get(url, timeout=10)
            time.sleep(self.REQUEST_DELAY)
            
            if response.status_code != 200:
                return None
            
            prop_data = response.json()['PropertyTable']['Properties'][0]
            
            # 创建化合物对象
            compound = CompoundData(
                cid=cid,
                molecular_formula=prop_data.get('MolecularFormula', ''),
                molecular_weight=float(prop_data.get('MolecularWeight', 0)),
                iupac_name=prop_data.get('IUPACName'),
                smiles=prop_data.get('CanonicalSMILES'),
                inchi=prop_data.get('InChI'),
                inchi_key=prop_data.get('InChIKey'),
                data_source="PubChem"
            )
            
            # 获取同义词（包括常用名称）
            try:
                syn_url = f"{self.PUBCHEM_BASE_URL}/compound/cid/{cid}/synonyms/JSON"
                syn_response = requests.get(syn_url, timeout=10)
                time.sleep(self.REQUEST_DELAY)
                
                if syn_response.status_code == 200:
                    synonyms = syn_response.json()['InformationList']['Information'][0]['Synonym']
                    compound.name = synonyms[0] if synonyms else f"CID_{cid}"
            except:
                compound.name = f"CID_{cid}"
            
            # 保存到缓存
            self.cache[cache_key] = compound.to_dict()
            self._save_cache()
            
            return compound
            
        except Exception as e:
            logger.error(f"查询失败 CID {cid}: {str(e)}")
            return None
    
    def batch_search(self, identifiers: List[str], search_type: str = "name") -> List[CompoundData]:
        """
        批量搜索化合物
        
        Args:
            identifiers: 标识符列表（名称或CAS号）
            search_type: 搜索类型 ("name" 或 "cas")
            
        Returns:
            化合物数据列表
        """
        results = []
        
        for identifier in identifiers:
            if search_type == "name":
                compound = self.search_by_name(identifier)
            elif search_type == "cas":
                compound = self.search_by_cas(identifier)
            else:
                logger.error(f"不支持的搜索类型: {search_type}")
                continue
            
            if compound:
                results.append(compound)
            
            # 避免请求过快
            time.sleep(self.REQUEST_DELAY)
        
        return results
    
    def export_to_json(self, compounds: List[CompoundData], output_file: str):
        """
        导出化合物数据到JSON文件
        
        Args:
            compounds: 化合物列表
            output_file: 输出文件路径
        """
        data = {
            "total": len(compounds),
            "compounds": [c.to_dict() for c in compounds]
        }
        
        output_path = Path(output_file)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        
        logger.info(f"已导出 {len(compounds)} 个化合物到 {output_file}")


# 使用示例
if __name__ == "__main__":
    # 创建数据库实例
    db = CompoundDatabase()
    
    # 示例1: 通过名称搜索
    print("=" * 80)
    print("示例1: 通过名称搜索化合物")
    print("=" * 80)
    
    ethanol = db.search_by_name("ethanol")
    if ethanol:
        print(f"名称: {ethanol.name}")
        print(f"分子式: {ethanol.molecular_formula}")
        print(f"分子量: {ethanol.molecular_weight}")
        print(f"SMILES: {ethanol.smiles}")
        print(f"InChI: {ethanol.inchi}")
    
    # 示例2: 通过CAS号搜索
    print("\n" + "=" * 80)
    print("示例2: 通过CAS号搜索")
    print("=" * 80)
    
    benzene = db.search_by_cas("71-43-2")
    if benzene:
        print(f"名称: {benzene.name}")
        print(f"CAS: {benzene.cas_number}")
        print(f"分子式: {benzene.molecular_formula}")
    
    # 示例3: 批量搜索
    print("\n" + "=" * 80)
    print("示例3: 批量搜索化合物")
    print("=" * 80)
    
    compounds = db.batch_search(["water", "methanol", "acetone"], search_type="name")
    for comp in compounds:
        print(f"{comp.name}: {comp.molecular_formula} (MW: {comp.molecular_weight})")

